package rbTree;

public class RBTree {
    public static class RBTreeNode {
        public RBTreeNode left;
        public RBTreeNode right;
        public RBTreeNode parent;
        public int val;
        public COLOR color;
        public RBTreeNode( int val){
            this.val = val;
            // 新创建的节点，默认是红色
            this.color = COLOR.RED;
        }
    }
    public RBTreeNode root;
    public boolean insert(int val) {
        RBTreeNode node = new RBTreeNode(val);
        if (root == null) {
            root = node;
            root.color = COLOR.BLACK;
            return true;
        }
        RBTreeNode cur = root;
        RBTreeNode parent = root;
        // 寻找插入的位置
        while (cur != null) {
            parent = cur;
            if (node.val < cur.val) {
                cur = cur.left;
            } else if (node.val == cur.val) {
                return false;
            } else {
                cur = cur.right;
            }
        }
        // 此时temp == null
        if (node.val < parent.val) {
            parent.left = node;
        } else {
            parent.right = node;
        }
        node.parent = parent;
        cur = node;
        // 插入进来之后需要调整颜色
        while (parent != null && parent.color == COLOR.RED) {
            RBTreeNode grandFather = parent.parent; //因为parent是红色 所以parent不是根节点
            if(grandFather == null){
                break;
            }
            if (parent == grandFather.left) {
                RBTreeNode uncle = grandFather.right;
                if (uncle != null && uncle.color == COLOR.RED) {
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;
                    cur = grandFather;
                    parent = cur.parent; // 继续向上修改
                } else {
                    if(cur == parent.right){
                        rotateLeft(parent);
                        cur = parent;
                        parent = cur.parent;
                    }
                    rotateRight(grandFather);
                    grandFather.color = COLOR.RED;
                    parent.color = COLOR.BLACK;
                }
            } else { // parent == grandFather.right
                RBTreeNode uncle = grandFather.left;
                if (uncle != null && uncle.color == COLOR.RED) {
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;
                    cur = grandFather;
                    parent = cur.parent;
                } else {
                    if(cur == parent.left){
                        rotateRight(parent);
                        cur = parent;
                        parent = cur.parent;
                    }
                    rotateLeft(grandFather);
                    grandFather.color = COLOR.RED;
                    parent.color = COLOR.BLACK;
                }
            }
        }
        root.color = COLOR.BLACK;
        return true;
    }
    private void rotateRight(RBTreeNode parent){
        RBTreeNode subL = parent.left;
        RBTreeNode subLR = subL.right;
        parent.left = subLR;
        subL.right = parent;
        if(subLR != null){
            subLR.parent = parent;
        }
        // 先记录parent的parent
        RBTreeNode pParent = parent.parent;
        parent.parent = subL;
        // 判断是否为根节点
        if(parent == root){
            root = subL;
            root.parent = null;
        } else { // 不是根节点时，判断是左子树还是右子树
            if(pParent.left == parent){
                pParent.left = subL;
            } else { // pParent.right = parent
                pParent.right = subL;
            }
            subL.parent = pParent;
        }
    }

    private void rotateLeft(RBTreeNode parent){
        RBTreeNode subR = parent.right;
        RBTreeNode subRL = subR.left;
        parent.right = subRL;
        subR.left = parent;
        if(subRL != null){
            subRL.parent = parent;
        }
        RBTreeNode pParent = parent.parent;
        parent.parent = subR;
        if (parent == root){
            root = subR;
            root.parent = null;
        } else {
            if(parent == pParent.left){
                pParent.left = subR;
            } else {
                pParent.right = subR;
            }
            subR.parent = pParent;
        }
    }

    /**
     * 判断一棵树是否是红黑树
     * 1. 根节点是黑色的
     * 2. 没有2个连续的红色节点
     * 3. 根节点到每一个叶子节点包含相同数目的黑色节点
     * @param root 根节点
     * @return boolean
     */
    public boolean isRBTree(RBTreeNode root){
        if(root == null){
            return true;
        }
        // 判断根节点
        if(root.color == COLOR.RED){
            System.out.println("根节点不是黑色的！");
            return false;
        }

        // 判断是否有两个连续的红色节点
        if(!checkRedNode(root)){
            return false;
        }

        // 计算根节点到最左边叶子结点路径中黑色节点的数量
        RBTreeNode cur = root;
        int countBlack = 0;
        while (cur != null){
            if(cur.color == COLOR.BLACK){
                countBlack++;
            }
            cur = cur.left;
        }

        // 判断其余路径的节点数量是否等于 countBlack
        if(!checkBlackCount(root, 0, countBlack)){
            return false;
        }
        return true;
    }
    private boolean checkRedNode(RBTreeNode root){
        if(root == null){
            return true;
        }
        if(root.color == COLOR.RED){
            RBTreeNode parent = root.parent;
            if(parent.color == COLOR.RED){
                System.out.println("有两个连续的红色节点！");
                return false;
            }
        }
        return checkRedNode(root.left) && checkRedNode(root.right);
    }

    private boolean checkBlackCount(RBTreeNode root, int countPath, int countBlack){
        if(root == null){
            return true;
        }
        if(root.color == COLOR.BLACK){
            countPath++;
        }
        if(root.left == null && root.right == null){
            if (countBlack != countPath){
                System.out.println("各个路径的黑色节点数量不相等");
                return false;
            }
            return true;
        }
        boolean leftTree = checkBlackCount(root.left, countPath, countBlack);
        boolean rightTree = checkBlackCount(root.right, countPath, countBlack);
        return leftTree && rightTree;
    }
}
